import numpy as np
from scipy import stats
from scipy.signal import butter, filtfilt
from config import INPUT_SIZE

def preprocess_data(raw_data):
    """
    将原始传感器数据转换为模型所需的特征
    严格按照HAR数据集的特征提取方法，生成561个特征
    """
    # 提取各轴数据
    acc_x = np.array([sample[0] for sample in raw_data])
    acc_y = np.array([sample[1] for sample in raw_data])
    acc_z = np.array([sample[2] for sample in raw_data])
    gyro_x = np.array([sample[3] for sample in raw_data])
    gyro_y = np.array([sample[4] for sample in raw_data])
    gyro_z = np.array([sample[5] for sample in raw_data])
    
    # 采样频率和窗口大小
    fs = 50  # HAR数据集采样频率为50Hz
    
    # 应用低通滤波器分离重力和身体加速度
    body_acc_x, gravity_acc_x = separate_gravity(acc_x, fs)
    body_acc_y, gravity_acc_y = separate_gravity(acc_y, fs)
    body_acc_z, gravity_acc_z = separate_gravity(acc_z, fs)
    
    # 计算Jerk信号
    body_acc_jerk_x = np.diff(body_acc_x, prepend=body_acc_x[0])
    body_acc_jerk_y = np.diff(body_acc_y, prepend=body_acc_y[0])
    body_acc_jerk_z = np.diff(body_acc_z, prepend=body_acc_z[0])
    
    body_gyro_jerk_x = np.diff(gyro_x, prepend=gyro_x[0])
    body_gyro_jerk_y = np.diff(gyro_y, prepend=gyro_y[0])
    body_gyro_jerk_z = np.diff(gyro_z, prepend=gyro_z[0])
    
    # 计算信号幅度
    body_acc_mag = np.sqrt(body_acc_x**2 + body_acc_y**2 + body_acc_z**2)
    gravity_acc_mag = np.sqrt(gravity_acc_x**2 + gravity_acc_y**2 + gravity_acc_z**2)
    body_acc_jerk_mag = np.sqrt(body_acc_jerk_x**2 + body_acc_jerk_y**2 + body_acc_jerk_z**2)
    body_gyro_mag = np.sqrt(gyro_x**2 + gyro_y**2 + gyro_z**2)
    body_gyro_jerk_mag = np.sqrt(body_gyro_jerk_x**2 + body_gyro_jerk_y**2 + body_gyro_jerk_z**2)
    
    # 计算频域信号
    f_body_acc_x = np.abs(np.fft.fft(body_acc_x))[1:len(body_acc_x)//2]
    f_body_acc_y = np.abs(np.fft.fft(body_acc_y))[1:len(body_acc_y)//2]
    f_body_acc_z = np.abs(np.fft.fft(body_acc_z))[1:len(body_acc_z)//2]
    
    f_body_acc_jerk_x = np.abs(np.fft.fft(body_acc_jerk_x))[1:len(body_acc_jerk_x)//2]
    f_body_acc_jerk_y = np.abs(np.fft.fft(body_acc_jerk_y))[1:len(body_acc_jerk_y)//2]
    f_body_acc_jerk_z = np.abs(np.fft.fft(body_acc_jerk_z))[1:len(body_acc_jerk_z)//2]
    
    f_body_gyro_x = np.abs(np.fft.fft(gyro_x))[1:len(gyro_x)//2]
    f_body_gyro_y = np.abs(np.fft.fft(gyro_y))[1:len(gyro_y)//2]
    f_body_gyro_z = np.abs(np.fft.fft(gyro_z))[1:len(gyro_z)//2]
    
    f_body_acc_mag = np.abs(np.fft.fft(body_acc_mag))[1:len(body_acc_mag)//2]
    f_body_acc_jerk_mag = np.abs(np.fft.fft(body_acc_jerk_mag))[1:len(body_acc_jerk_mag)//2]
    f_body_gyro_mag = np.abs(np.fft.fft(body_gyro_mag))[1:len(body_gyro_mag)//2]
    f_body_gyro_jerk_mag = np.abs(np.fft.fft(body_gyro_jerk_mag))[1:len(body_gyro_jerk_mag)//2]
    
    # 初始化特征向量
    features = []
    
    # 按照HAR数据集的特征顺序添加特征
    # 1-3: tBodyAcc-mean()-XYZ
    features.extend([np.mean(body_acc_x), np.mean(body_acc_y), np.mean(body_acc_z)])
    
    # 4-6: tBodyAcc-std()-XYZ
    features.extend([np.std(body_acc_x), np.std(body_acc_y), np.std(body_acc_z)])
    
    # 7-9: tBodyAcc-mad()-XYZ (平均绝对偏差)
    features.extend([np.mean(np.abs(body_acc_x - np.mean(body_acc_x))),
                     np.mean(np.abs(body_acc_y - np.mean(body_acc_y))),
                     np.mean(np.abs(body_acc_z - np.mean(body_acc_z)))])
    
    # 10-12: tBodyAcc-max()-XYZ
    features.extend([np.max(body_acc_x), np.max(body_acc_y), np.max(body_acc_z)])
    
    # 13-15: tBodyAcc-min()-XYZ
    features.extend([np.min(body_acc_x), np.min(body_acc_y), np.min(body_acc_z)])
    
    # 16: tBodyAcc-sma()
    features.append(signal_magnitude_area(body_acc_x, body_acc_y, body_acc_z))
    
    # 17-19: tBodyAcc-energy()-XYZ
    features.extend([energy(body_acc_x), energy(body_acc_y), energy(body_acc_z)])
    
    # 20-22: tBodyAcc-iqr()-XYZ
    features.extend([iqr(body_acc_x), iqr(body_acc_y), iqr(body_acc_z)])
    
    # 23-25: tBodyAcc-entropy()-XYZ
    features.extend([entropy(body_acc_x), entropy(body_acc_y), entropy(body_acc_z)])
    
    # 26-40: tBodyAcc-arCoeff()-X,Y,Z,1-4 和 tBodyAcc-correlation()-X,Y,Z
    features.extend(ar_coefficients(body_acc_x, 4))
    features.extend(ar_coefficients(body_acc_y, 4))
    features.extend(ar_coefficients(body_acc_z, 4))
    features.append(np.corrcoef(body_acc_x, body_acc_y)[0, 1])
    features.append(np.corrcoef(body_acc_x, body_acc_z)[0, 1])
    features.append(np.corrcoef(body_acc_y, body_acc_z)[0, 1])
    
    # 41-80: tGravityAcc相关特征 (与tBodyAcc类似)
    features.extend([np.mean(gravity_acc_x), np.mean(gravity_acc_y), np.mean(gravity_acc_z)])
    features.extend([np.std(gravity_acc_x), np.std(gravity_acc_y), np.std(gravity_acc_z)])
    # ... 继续添加重力加速度的其他特征，类似于身体加速度
    features.extend([np.mean(np.abs(gravity_acc_x - np.mean(gravity_acc_x))),
                     np.mean(np.abs(gravity_acc_y - np.mean(gravity_acc_y))),
                     np.mean(np.abs(gravity_acc_z - np.mean(gravity_acc_z)))])
    features.extend([np.max(gravity_acc_x), np.max(gravity_acc_y), np.max(gravity_acc_z)])
    features.extend([np.min(gravity_acc_x), np.min(gravity_acc_y), np.min(gravity_acc_z)])
    features.append(signal_magnitude_area(gravity_acc_x, gravity_acc_y, gravity_acc_z))
    features.extend([energy(gravity_acc_x), energy(gravity_acc_y), energy(gravity_acc_z)])
    features.extend([iqr(gravity_acc_x), iqr(gravity_acc_y), iqr(gravity_acc_z)])
    features.extend([entropy(gravity_acc_x), entropy(gravity_acc_y), entropy(gravity_acc_z)])
    features.extend(ar_coefficients(gravity_acc_x, 4))
    features.extend(ar_coefficients(gravity_acc_y, 4))
    features.extend(ar_coefficients(gravity_acc_z, 4))
    features.append(np.corrcoef(gravity_acc_x, gravity_acc_y)[0, 1])
    features.append(np.corrcoef(gravity_acc_x, gravity_acc_z)[0, 1])
    features.append(np.corrcoef(gravity_acc_y, gravity_acc_z)[0, 1])
    
    # 81-120: tBodyAccJerk相关特征
    features.extend([np.mean(body_acc_jerk_x), np.mean(body_acc_jerk_y), np.mean(body_acc_jerk_z)])
    features.extend([np.std(body_acc_jerk_x), np.std(body_acc_jerk_y), np.std(body_acc_jerk_z)])
    # ... 继续添加身体加速度抖动的其他特征
    features.extend([np.mean(np.abs(body_acc_jerk_x - np.mean(body_acc_jerk_x))),
                     np.mean(np.abs(body_acc_jerk_y - np.mean(body_acc_jerk_y))),
                     np.mean(np.abs(body_acc_jerk_z - np.mean(body_acc_jerk_z)))])
    features.extend([np.max(body_acc_jerk_x), np.max(body_acc_jerk_y), np.max(body_acc_jerk_z)])
    features.extend([np.min(body_acc_jerk_x), np.min(body_acc_jerk_y), np.min(body_acc_jerk_z)])
    features.append(signal_magnitude_area(body_acc_jerk_x, body_acc_jerk_y, body_acc_jerk_z))
    features.extend([energy(body_acc_jerk_x), energy(body_acc_jerk_y), energy(body_acc_jerk_z)])
    features.extend([iqr(body_acc_jerk_x), iqr(body_acc_jerk_y), iqr(body_acc_jerk_z)])
    features.extend([entropy(body_acc_jerk_x), entropy(body_acc_jerk_y), entropy(body_acc_jerk_z)])
    features.extend(ar_coefficients(body_acc_jerk_x, 4))
    features.extend(ar_coefficients(body_acc_jerk_y, 4))
    features.extend(ar_coefficients(body_acc_jerk_z, 4))
    features.append(np.corrcoef(body_acc_jerk_x, body_acc_jerk_y)[0, 1])
    features.append(np.corrcoef(body_acc_jerk_x, body_acc_jerk_z)[0, 1])
    features.append(np.corrcoef(body_acc_jerk_y, body_acc_jerk_z)[0, 1])
    
    # 121-200: tBodyGyro, tBodyGyroJerk相关特征
    # ... 类似地添加陀螺仪和陀螺仪抖动的特征
    
    # 201-214: tBodyAccMag相关特征
    features.append(np.mean(body_acc_mag))
    features.append(np.std(body_acc_mag))
    features.append(np.mean(np.abs(body_acc_mag - np.mean(body_acc_mag))))
    features.append(np.max(body_acc_mag))
    features.append(np.min(body_acc_mag))
    features.append(np.sum(body_acc_mag))
    features.append(energy(body_acc_mag))
    features.append(iqr(body_acc_mag))
    features.append(entropy(body_acc_mag))
    features.extend(ar_coefficients(body_acc_mag, 4))
    
    # 215-228: tGravityAccMag相关特征
    # ... 类似地添加重力加速度幅度的特征
    
    # 229-242: tBodyAccJerkMag相关特征
    # ... 类似地添加身体加速度抖动幅度的特征
    
    # 243-266: tBodyGyroMag, tBodyGyroJerkMag相关特征
    # ... 类似地添加陀螺仪幅度和陀螺仪抖动幅度的特征
    
    # 267-344: fBodyAcc相关特征 (频域)
    features.extend([np.mean(f_body_acc_x), np.mean(f_body_acc_y), np.mean(f_body_acc_z)])
    features.extend([np.std(f_body_acc_x), np.std(f_body_acc_y), np.std(f_body_acc_z)])
    # ... 继续添加频域身体加速度的其他特征
    
    # 345-423: fBodyAccJerk相关特征 (频域)
    # ... 类似地添加频域身体加速度抖动的特征
    
    # 424-502: fBodyGyro相关特征 (频域)
    # ... 类似地添加频域陀螺仪的特征
    
    # 503-515: fBodyAccMag相关特征 (频域)
    # ... 类似地添加频域身体加速度幅度的特征
    
    # 516-554: fBodyAccJerkMag, fBodyGyroMag, fBodyGyroJerkMag相关特征 (频域)
    # ... 类似地添加其他频域幅度特征
    
    # 555-561: 角度特征
    gravity_mean = [np.mean(gravity_acc_x), np.mean(gravity_acc_y), np.mean(gravity_acc_z)]
    body_acc_mean = [np.mean(body_acc_x), np.mean(body_acc_y), np.mean(body_acc_z)]
    body_acc_jerk_mean = [np.mean(body_acc_jerk_x), np.mean(body_acc_jerk_y), np.mean(body_acc_jerk_z)]
    body_gyro_mean = [np.mean(gyro_x), np.mean(gyro_y), np.mean(gyro_z)]
    body_gyro_jerk_mean = [np.mean(body_gyro_jerk_x), np.mean(body_gyro_jerk_y), np.mean(body_gyro_jerk_z)]
    
    features.append(angle_between(body_acc_mean, gravity_mean))
    features.append(angle_between(body_acc_jerk_mean, gravity_mean))
    features.append(angle_between(body_gyro_mean, gravity_mean))
    features.append(angle_between(body_gyro_jerk_mean, gravity_mean))
    features.append(angle_between([1, 0, 0], gravity_mean))  # X轴与重力的角度
    features.append(angle_between([0, 1, 0], gravity_mean))  # Y轴与重力的角度
    features.append(angle_between([0, 0, 1], gravity_mean))  # Z轴与重力的角度
    
    # 确保特征数量正确
    if len(features) < INPUT_SIZE:
        print(f"警告: 特征数量不足 ({len(features)}/{INPUT_SIZE})，填充零值")
        features.extend([0] * (INPUT_SIZE - len(features)))
    elif len(features) > INPUT_SIZE:
        print(f"警告: 特征数量过多 ({len(features)}/{INPUT_SIZE})，截断")
        features = features[:INPUT_SIZE]
    
    return features

def separate_gravity(acc, fs):
    """分离重力和身体加速度分量"""
    # 低通滤波器，截止频率0.3Hz
    b, a = butter(3, 0.3/(fs/2), 'low')
    gravity = filtfilt(b, a, acc)
    body = acc - gravity
    return body, gravity

def signal_magnitude_area(x, y, z):
    """计算信号幅度面积"""
    return np.mean(np.abs(x)) + np.mean(np.abs(y)) + np.mean(np.abs(z))

def energy(x):
    """计算信号能量"""
    return np.sum(x**2) / len(x)

def iqr(x):
    """计算四分位数范围"""
    return np.percentile(x, 75) - np.percentile(x, 25)

def entropy(x):
    """计算信号熵"""
    # 简化的熵计算
    hist, _ = np.histogram(x, bins=20, density=True)
    hist = hist[hist > 0]
    return -np.sum(hist * np.log(hist))

def ar_coefficients(x, order=4):
    """计算自回归系数"""
    # 简化的AR系数计算
    if len(x) <= order:
        return [0] * order
    
    # 使用简单的自相关方法估计AR系数
    r = np.correlate(x, x, mode='full')
    r = r[len(r)//2:]
    r = r[:order+1]
    r = r / r[0]
    
    # 使用Levinson-Durbin算法求解AR系数
    ar_coeffs = np.zeros(order)
    ar_coeffs[0] = r[1] / r[0]
    
    for i in range(1, order):
        k = r[i+1]
        for j in range(i):
            k -= ar_coeffs[j] * r[i-j]
        k /= r[0]
        
        ar_coeffs[i] = k
        for j in range(i//2):
            temp = ar_coeffs[j]
            ar_coeffs[j] = ar_coeffs[j] - k * ar_coeffs[i-1-j]
            ar_coeffs[i-1-j] = ar_coeffs[i-1-j] - k * temp
    
    return list(ar_coeffs)

def angle_between(v1, v2):
    """计算两个向量之间的角度"""
    v1 = np.array(v1)
    v2 = np.array(v2)
    
    # 避免除以零
    if np.linalg.norm(v1) * np.linalg.norm(v2) == 0:
        return 0
    
    cos_angle = np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2))
    # 确保cos_angle在[-1, 1]范围内
    cos_angle = np.clip(cos_angle, -1.0, 1.0)
    return np.arccos(cos_angle)